# 同时采集像素和对应车辆坐标
import os
import sys
import cv2
import numpy as np
import math
import rospy
from nav_msgs.msg import Odometry
import time
from collections import deque
from typing import Union, List
import threading

os.chdir(sys.path[0])  # 便于在vscode直接运行
import math
import cv2 as cv
import numpy as np
from glob import glob


class RosNode:
    def __init__(self, node_name='navigation_node', debug_mode=False):
        self.node_name = node_name
        self.carla_odometry_topic = '/carla/ego_vehicle/odometry'
        self.is_odometry_received = False
        self.carla_odometry_msg = Odometry()
        self.msg_pool = {}
        self.debug_mode = debug_mode

    def ros_init(self):
        rospy.init_node(self.node_name, anonymous=True, disable_signals=True)
        self.pool_regist(['odometry'])  # 注册消息池
        self.listener()

    def listener(self):
        rospy.Subscriber(
            self.carla_odometry_topic,
            Odometry,
            callback=self.callback_odometry,
            queue_size=1,
        )

    def callback_odometry(self, data: Odometry):
        this_time = rospy.get_time() * 1000.0  # ms
        self.is_odometry_received = True
        self.carla_odometry_msg = data
        self.msg_pool_append("odometry", data, this_time)
        if self.debug_mode:
            print(f"get {type(data)} msg")

    def msg_pool_append(self, msg_name: str, msg_data, msg_time: float):
        """消息池存储函数

        Args:
            msg_name (str): _description_
            msg_data (_type_): 消息data
            msg_time (rostime): 消息对应的ros时间
        """
        msg_tuple = (msg_data, msg_time)
        try:
            self.msg_pool[msg_name].append(msg_tuple)
        except KeyError:
            self.msg_pool[msg_name].append(msg_tuple)

    def pool_regist(self, msg_name: Union[str, List[str]]):
        """消息池中字典注册

        Args:
            msg_name (Union[str, List[str]]): 消息名，或者包含消息名的列表
        """
        if isinstance(msg_name, list):
            for m in msg_name:
                self.pool_regist(m)
        elif isinstance(msg_name, str):
            self.msg_pool[msg_name] = deque(maxlen=5)


def readImage(dir, idex):
    """读取图片"""
    img_lists = glob(dir + '*')
    img_lists.sort()
    img1 = cv.imread(img_lists[idex])  # 这里必须用cv库里面的imread，否则格式不对会报错
    return img1


# 定义一个窗口
def draw_(event, x, y, flag, param):
    global data, img1, loc_carla, loc_carla_list
    if event == cv.EVENT_LBUTTONDOWN:
        """左键单击打标签"""
        xy = "%d,%d" % (x, y)
        print(f"u:{x},v:{y}\n")
        print(f"carla loc :{loc_carla}\n")
        cv.circle(img1, (x, y), 2, (255, 0, 0), thickness=-1)
        data.append([x, y])
        loc_carla_list.append([loc_carla[0][0], loc_carla[1][0]])

    if event == cv.EVENT_MBUTTONDOWN:
        """单击鼠标中间，存储数据"""
        file = open('input_image_points.txt', mode='w+')
        for i in range(len(data)):
            u_str = str(data[i][0])
            v_str = str(data[i][1])
            file.writelines(u_str + ' ' + v_str + '\n')
        file.close()
        print("file saved to :input_image_points.txt")

        file = open('input_world_points.txt', mode='w+')
        for i in range(len(loc_carla_list)):
            u_str = str(loc_carla_list[i][0])
            v_str = str(loc_carla_list[i][1])
            file.writelines(u_str + ' ' + v_str + '\n')
        file.close()
        print("file saved to :input_world_points.txt")


class imgThread(threading.Thread):
    def __init__(self, threadId, name):
        threading.Thread.__init__(self)
        self.threadId = threadId
        self.name = name
        self.counter = 0

    def run(self):
        global img1
        cv.namedWindow('image')
        cv.setMouseCallback('image', draw_)
        print("开始线程:", self.name)
        while 1:
            # 需要运行的程序：
            cv.imshow('image', img1)
            k = cv.waitKey(1) & 0xFF
            if k == ord('c'):  # 按c退出
                break
            time.sleep(1e-2)  # 睡眠10ms
        cv.destroyAllWindows()

    def __del__(self):
        print(self.name, "线程结束！")


class locThread(threading.Thread):
    def __init__(self, threadId, name):
        threading.Thread.__init__(self)
        self.threadId = threadId
        self.name = name
        self.counter = 0

    def run(self):
        global img1, loc_carla, loc_thread
        print("开始线程:", self.name)
        rate = rospy.Rate(5)
        while not rospy.is_shutdown() and loc_thread.is_alive():
            # 需要运行的程序：
            if my_node.is_odometry_received:
                # 从ros中取消息
                carla_msg = my_node.msg_pool['odometry'][-1][0]
                loc_x = carla_msg.pose.pose.position.x
                loc_y = carla_msg.pose.pose.position.y
                loc_carla = np.array([[loc_x], [loc_y]])
                # print(loc_carla)

            rate.sleep()

    def __del__(self):
        print(self.name, "线程结束！")


if __name__ == "__main__":
    idex = 6
    # 读取图片
    dir = '/home/wrc/study/code/carla_tutorial/transform/transferMatric_Loc2Ima/get_pixel/pic/'  # 图片存放的目录

    # 创建一个列表存储标签和RGB值
    data = []
    label = []
    img1 = readImage(dir, idex)

    # 存储carla坐标
    loc_carla = None
    loc_carla_list = []

    # 启动主程序
    my_node = RosNode('carla_listen_node', debug_mode=False)
    my_node.ros_init()

    # 启动数据采集线程
    loc_thread = locThread(1, 'msgThread')
    loc_thread.start()

    # 启动地图绘制线程
    img_thread = imgThread(2, 'imgThread')
    img_thread.start()

    rospy.spin()
